#include <cassert>
#include <iostream>

using namespace std;

#include "type.hpp"
#include "ast.hpp"

bool IntType::support_operator(OperatorType ot, Type* other) {
    if (ot == OPERATOR_ADD || ot == OPERATOR_SUB ||
            ot == OPERATOR_MUL || ot == OPERATOR_DIV ||
            ot == OPERATOR_SHIFT) {
        //return unify(other, IntType::get_instance());
        return true;
    }

    return false;
}

bool CharType::support_operator(OperatorType ot, Type* other) {
    if (ot == OPERATOR_ADD || ot == OPERATOR_SUB) {
        //return unify(other, CharType::get_instance());
        return true;
    }

    return false;
}

bool StringType::support_operator(OperatorType ot, Type* other) {
    if (ot == OPERATOR_ADD) {
        //return unify(other, StringType::get_instance());
        return true;
    }
    else if (ot == OPERATOR_MUL) {
        //return unify(other, IntType::get_instance());
        return true;
    }

    return false;
}

bool DoubleType::support_operator(OperatorType ot, Type* other) {
    if (ot == OPERATOR_ADD || ot == OPERATOR_SUB ||
            ot == OPERATOR_MUL || ot == OPERATOR_DIV ||
            ot == OPERATOR_MOD) {
        //return unify(other, DoubleType::get_instance());
        return true;
    }

    return false;
}

bool ListType::support_operator(OperatorType ot, Type* other) {
    if (ot == OPERATOR_ADD) {
        // return unify(other, ListType::get_instance());
        return true;
    }
    else if (ot == OPERATOR_MUL) {
        //return unify(other, IntType::get_instance());
        return true;
    }

    return false;
}

ArrowType::~ArrowType() {
}

string ArrowType::to_string() {
    return _src->to_string() + " -> " + _dst->to_string();
}

bool ArrowType::equals(Type* t) {
    ArrowType* at = dynamic_cast<ArrowType*>(t);
    if (at == NULL) {
        return false;
    }

    return _src->equals(at->src()) && _dst->equals(at->dst());
}

Type* ArrowType::try_apply(Type* t) {
    if (!t) {
        cerr << "ArrowType::try_apply : type is nullptr." << endl;
    }
    return _dst;
}

Type* ArrowType::substitute(TypeTable* type_table) {
    if (is_substing()) {
        return this;
    }
    set_substing(true);
    Type* src = _src->substitute(type_table);
    Type* dst = _dst->substitute(type_table);
    set_substing(false);
    return new ArrowType(src, dst);
}

Type* TupleType::substitute(TypeTable* type_table) {
    vector<Type*> types;
    for (auto it = _types.begin(); it != _types.end(); it++) {
        types.push_back((*it)->substitute(type_table));
    }
    return new TupleType(types);
}

bool TupleType::equals(Type* t) {
    TupleType* tt = dynamic_cast<TupleType*>(t);
    if (tt == NULL) {
        return false;
    }

    if (_types.size() != tt->_types.size()) {
        return false;
    }

    for (int i = 0; i < _types.size(); i++) {
        if (!_types[i]->equals(tt->_types[i])) {
            return false;
        }
    }

    return true;
}

string TupleType::to_string() {
    string s = "(";
    for (int i = 0; i < _types.size(); i++) {
        s += _types[i]->to_string();
        if (i != _types.size() - 1) {
            s += ", ";
        }
    }
    s += ")";
    return s;
}

Type* TypeVar::substitute(TypeTable* type_table) {
    return env()->ufset_find(this);
}

bool TypeVar::equals(Type* t) {
    TypeVar* tv = dynamic_cast<TypeVar*>(t);
    if (tv == nullptr) {
        return false;
    }
    return name() == tv->name();
}

Type* RecordType::substitute(TypeTable* type_table) {
    if (is_substing()) {
        return this;
    }
    set_substing(true);

    map<string, Type*> fields;
    vector<Type*> field_types;
    map<string, Type*> methods;
    for (auto it = _fields.begin(); it != _fields.end(); it++) {
        fields[it->first] = it->second->substitute(type_table);
        field_types.push_back(fields[it->first]);
    }

    for (auto it = _method_types.begin(); it != _method_types.end(); it++) {
        methods[it->first] = it->second->substitute(type_table);
    }

    RecordType* rt = new RecordType(_name, fields, field_types);
    rt->set_methods_types(methods);
    set_substing(false);
    return rt;
}

bool RecordType::support_operator(OperatorType ot, Type* other) {

    if (ot == OPERATOR_ADD && has_method("__add__")) {
        return true;
    }
    else if (ot == OPERATOR_SUB && has_method("__sub__")) {
        return true;
    }
    else if (ot == OPERATOR_MUL && has_method("__mul__")) {
        return true;
    }
    else if (ot == OPERATOR_DIV && has_method("__div__")) {
        return true;
    }
    else if (ot == OPERATOR_MOD && has_method("__mod__")) {
        return true;
    }
    else if (ot == OPERATOR_SHIFT && has_method("__shift__")) {
        return true;
    }
    return false;
}

Type* TypeTable::get_substitution(string& name) {
    auto it = _substitution_map.find(name);
    if (it != _substitution_map.end()) {
        return it->second;
    }
    else if (parent()) {
        return parent()->get_substitution(name);
    }
    return nullptr;
}

// union find set find.
Type* TypeTable::ufset_find(Type* t) {
    if (!t->is_type_var()) {
        return t;
    }

    TypeVar* tv = dynamic_cast<TypeVar*>(t);
    Type* sbt = tv;
    while (_substitution_map.find(tv->name()) != _substitution_map.end()) {
        sbt = _substitution_map[tv->name()];
        if (sbt->is_type_var()) {
            tv = dynamic_cast<TypeVar*>(sbt);
            sbt = tv;
        }
        else {
            return sbt;
        }
    }
    
    if (_parent) {
        return _parent->ufset_find(sbt);
    }
    return sbt;
}

void TypeTable::ufset_union(Type* t1, Type* t2) {
    assert(t1->is_type_var());
    TypeVar* tv1 = dynamic_cast<TypeVar*>(t1);
    tv1->env()->add_substitution(tv1->name(), t2);
}

Type* TypeVar::get_type(TypeTable* type_table) {
    return env()->ufset_find(this);
}

Type* PolyType::expand(TypeTable* type_table) {
    return type_constructor->substitute(type_table);
}

Type* TypeFunction::expand_trait(TypeTable* type_table) {
    assert(false && "TypeFunction::expand_trait : not implemented.");
    return _body->substitute(type_table);
}

Type* TypeFunction::expand_record(TypeTable* type_table) {
    RecordType* rt = dynamic_cast<RecordType*>(_body);
    return rt->substitute(type_table);
}

Type* TypeFunction::expand(TypeTable* type_table) {
    if (_body->is_record()) {
        return expand_record(type_table);
    }
    else if (_body->is_arrow()) {
        return _body->substitute(type_table);
    }
    return _body;
}

Type* TypeApply::expand(TypeTable* type_table) {
    return type_constructor->substitute(type_table);
}

Type* ListType::substitute(TypeTable* type_table) {
    return new ListType(_elm_type->substitute(type_table));
}

bool ListType::equals(Type* t) {
    ListType* at = dynamic_cast<ListType*>(t);
    if (at == NULL) {
        return false;
    }

    return _elm_type->equals(at->elm_type());
}

UserDefType::~UserDefType() {
}

TypeTable::TypeTable() {
    _type_map["Int"] = IntType::get_instance();
    _type_map["Double"] = DoubleType::get_instance();
    _type_map["String"] = StringType::get_instance();
    _type_map["Char"] = CharType::get_instance();
    _type_map["Bool"] = BoolType::get_instance();
    _type_map["Unit"] = UnitType::get_instance();
}

TypeTable::TypeTable(TypeTable* parent) {
    _parent = parent;
}

void TypeTable::add_generic_name(Node* n) {
    TypeVarNode* tvn = dynamic_cast<TypeVarNode*>(n);
    assert(tvn != nullptr);

    _generic_names.push_back(tvn->name());
    _substitution_map[tvn->name()] = new TypeVar(tvn->name(), this);
}

Type* TypeTable::get_expanded_type(string& name) {
    TypeTable* tt = this;
    while (tt != nullptr) {
        auto t = tt->_type_map.find(name);
        if (t != tt->_type_map.end()) {
            return t->second;
        }
        tt = tt->parent();
    }

    return nullptr;
}

void TypeTable::clear_instantiation() {
    _substitution_map.clear();
}
